import random                     
import math                       
from typing import List, Tuple    # 从 typing 中导入类型提示，便于代码可读性

# 1. 计算欧氏距离的平方（避免不必要的 sqrt）
def euclidean_sq(a: List[float], b: List[float]) -> float:
    # 对两个向量对应元素求差的平方并求和，得到距离的平方
    return sum((x - y) ** 2 for x, y in zip(a, b))

# 2. 为单个样本分配最近的簇
def assign_cluster(x: List[float], centroids: List[List[float]]) -> int:
    """
    返回样本 x 最近的质心索引
    """
    distances = [euclidean_sq(x, c) for c in centroids]# 计算样本 x 与每个质心的距离平方，得到一个列表
    return distances.index(min(distances))# 返回距离最小的质心对应的索引

# 3. 根据当前标签重新计算质心
def compute_centroids(
    data: List[List[float]],   
    labels: List[int],        
    k: int                     
) -> List[List[float]]:
    """
    data   : 所有样本
    labels : 每个样本对应的簇编号
    k      : 簇的数量
    返回   : k 个新质心（列表形式）
    """
    dim = len(data[0])  # 取特征维度（假设所有样本维度相同）   
    centroids = [[0.0] * dim for _ in range(k)]# 初始化 k 个质心，每个质心是一个全 0 的向量
    counts = [0] * k# 用于统计每个簇中样本的数量

    # 累加每个簇中所有样本的特征值
    for x, label in zip(data, labels):
        counts[label] += 1                  # 该簇样本计数加一
        for i in range(dim):
            centroids[label][i] += x[i]     # 对应维度的特征值累加

    # 处理可能出现的空簇（即某个簇没有分配到任何样本）
    for idx in range(k):
        if counts[idx] == 0:
            # 若空簇，随机挑选一个样本作为该簇的新质心，防止除以 0
            centroids[idx] = random.choice(data)
        else:
            # 否则，用累加的特征和除以样本数，得到均值即新质心
            centroids[idx] = [v / counts[idx] for v in centroids[idx]]

    return centroids

# 4. K‑means 主函数
def kmeans(
    data: List[List[float]],   
    k: int,                   
    epsilon: float = 1e-4,     
    max_iter: int = 300        
) -> Tuple[List[List[float]], List[int]]:
    """
    data      : 输入数据，形如 [[x1, x2, ...], [y1, y2, ...], ...]
    k         : 目标簇数
    epsilon   : 质心移动阈值，低于该值视为收敛
    max_iter  : 最大迭代次数
    返回值    : (final_centroids, final_labels)
    """
    # 参数合法性检查
    if k <= 0:
        raise ValueError("k 必须为正整数")
    if k > len(data):
        raise ValueError("k 不能大于样本数")

    # 1) 随机初始化质心（从样本中随机抽取 k 个不重复的点）
    centroids = random.sample(data, k)

    # 迭代过程
    for it in range(max_iter):
        # 2) 为每个样本分配最近的簇
        labels = [assign_cluster(x, centroids) for x in data]

        # 3) 根据当前标签重新计算质心
        new_centroids = compute_centroids(data, labels, k)

        # 4) 检查收敛：计算所有质心位移的平方和的最大值
        max_shift = max(euclidean_sq(c, nc) for c, nc in zip(centroids, new_centroids))
        centroids = new_centroids   # 更新质心为新质心

        # 打印本次迭代的信息
        print(f"Iteration {it+1}: max centroid shift = {math.sqrt(max_shift):.6f}")

        # 若最大位移小于阈值的平方，则认为已收敛
        if max_shift < epsilon ** 2:
            print("收敛！")
            break

    # 返回最终的质心和每个样本的簇标签
    return centroids, labels

# 5. 示例（可直接运行）
if __name__ == "__main__":
    # 生成一个简单的二维数据集（手写，不依赖外部库）
    data = [
        [1.0, 2.0], [1.2, 1.8], [0.8, 2.2],   # 簇 0 的样本
        [5.0, 8.0], [5.2, 7.9], [4.9, 8.1],   # 簇 1 的样本
        [9.0, 1.0], [9.2, 1.1], [8.8, 0.9]    # 簇 2 的样本
    ]

    k = 3                                 # 设定要划分的簇数为 3
    # 调用 kmeans，epsilon 设置得更小以获得更精细的收敛
    final_centroids, final_labels = kmeans(data, k, epsilon=1e-5, max_iter=100)

    print("\n最终质心：")
    for idx, c in enumerate(final_centroids):
        print(f"  簇 {idx}: {c}")

    print("\n样本所属簇：")
    for idx, label in enumerate(final_labels):
        print(f"  样本 {idx} -> 簇 {label}")
